import asyncio
import time

from aioredis import Redis


# 自旋加锁
async def add_lock_self(redis: Redis, key, value, expire_time: int, execute_time: int):
    """

    :param redis:
    :param key:
    :param value:
    :param expire_time: 锁过期时间
    :param execute_time: 执行时间
    :return:

    """
    start_time = time.time() * 1000  # 单位: ms
    while True:
        await asyncio.sleep(expire_time / 1000 / 3)  # 三分之一的过期时间刷新一下
        lua1 = """
            if (redis.call('exists', KEYS[1]) == 0)
            then
                return -1
            else
                if (redis.call("get", KEYS[1]) == ARGV[1])
                then
                    redis.call("PEXPIRE", KEYS[1], ARGV[2])
                    return 1
                else
                    return 0
                end
            end
        """
        res = await redis.eval(lua1, 1, key, value, expire_time)
        if res == 0 or res == -1:  # 加锁失败 直接返回
            break
        # 加锁成功
        e = time.time() * 1000
        if e - start_time > execute_time:  # 如果超出了执行时间，直接返回
            break


# 加锁
async def add_lock(redis: Redis, key, value, expire_time: int, execute_time: int):
    start_time = time.time() * 1000
    while True:
        await asyncio.sleep(expire_time / 1000 / 10)  # 每隔十分之一的过期时间加锁一次
        lua = """
            if (redis.call('exists', KEYS[1]) == 0)
            then
                redis.call("set", KEYS[1], ARGV[1])
                redis.call("PEXPIRE", KEYS[1], ARGV[2])
                return 1
            else
                return 0
            end
        """
        res = await redis.eval(lua, 1, key, value, expire_time)
        if res:  # 加锁成功
            return True
        if time.time() * 1000 - start_time > execute_time:  # 超出执行时间， 加锁失败
            return False


async def release_lock(redis: Redis, key, request_id):
    rid = await redis.get(key)
    if rid and rid.decode() == request_id:
        await redis.delete(key)
        return True
    return False


# 添加分布式锁
async def add_red_lock_once(redis: Redis, key, value, expire_time):
    lua = """
        local ok = redis.call('setnx', KEYS[1], ARGV[1])
        if (ok == 1)
        then
            redis.call("EXPIRE", KEYS[1], ARGV[2])
        end
        return ok
    """
    # G_REDIS.set(k, v, nx=True, ex=)
    return await redis.eval(lua, 1, key, value, expire_time)
